from .base import PoseEstimator, center_ellipse, ransac_pnp, pose_calculate
import numpy as np
import cv2
from scipy.optimize import least_squares


class ED(PoseEstimator):
    def __init__(self, K, use_pnp_initial=True, *args, **kwargs):
        super().__init__(K, *args, **kwargs)
        self.use_pnp_initial = use_pnp_initial

    def __name__(self):
        return "ed"

    def forward(self, C_3d, C_2d, *args,**ignore_kwargs):
        x_cnt = []
        X_cnt = []
        for c_3d, c_2d in zip(C_3d, C_2d):
            # 求解相机椭圆中心：
            x_cnt.append(center_ellipse(c_2d))
            X_cnt.append(center_ellipse(c_3d))
        # 求解相机位姿
        X_cnt = np.pad(X_cnt, ((0, 0), (0, 1)), "constant", constant_values=0)
        x_cnt = np.array(x_cnt)
        if self.use_pnp_initial:
            ind = ransac_pnp(self.K, X_cnt, x_cnt)
            if ind is None:
                return False, None, None
            R, T = pose_calculate(self.K, X_cnt[ind], x_cnt[ind])
            initial_pose = np.concatenate((cv2.Rodrigues(R)[0], T)).flatten()
        else:
            initial_pose = np.zeros(6)
        return self.estimator(X_cnt, x_cnt, initial_pose)

    def estimator(self, database, measurements, initial_pose):
        """
        K: 相机内参
        database: 三维点库
        measurements: 观测到的二维点
        """
        assert len(database) == len(measurements)
        assert measurements.shape[1] == 2
        assert database.shape[1] == 3
        # 求解相机位姿

        def residual_function(params, observations, points_3d):
            # 定义残差函数
            R_ = cv2.Rodrigues(params[:3])[0]
            T_ = params[3:]
            uv = cv2.projectPoints(points_3d, R_, T_, self.K, None)[0].squeeze()
            return (uv - observations).flatten()

        # def jac(params, observations, points_3d):
        #     # 定义雅克比矩阵
        #     J = np.zeros((observations.shape[0] * 2, 6))
        #     R_ = cv2.Rodrigues(params[:3])[0].flatten()
        #     T_ = params[3:]
        #     for i in range(observations.shape[0]):
        #         x, y, z = points_3d[i]
        #         X = R_[0] * x + R_[1] * y + R_[2] * z + T_[0]
        #         Y = R_[3] * x + R_[4] * y + R_[5] * z + T_[1]
        #         Z = R_[6] * x + R_[7] * y + R_[8] * z + T_[2]
        #         Z_2 = Z * Z
        #         J[i * 2, :3] = (
        #             -1.0 / Z * np.array([[self.K[0, 0], 0, -X / Z * self.K[0, 0]]])
        #         )
        #         J[i * 2 + 1, :3] = (
        #             -1.0 / Z * np.array([[0, self.K[1, 1], -Y / Z * self.K[1, 1]]])
        #         )
        #         J[i * 2, 3] = X * Y / Z_2 * self.K[0, 0]
        #         J[i * 2, 4] = -(1 + (X * X / Z_2)) * self.K[0, 0]
        #         J[i * 2, 5] = Y / Z * self.K[0, 0]
        #         J[i * 2 + 1, 3] = (1 + Y * Y / Z_2) * self.K[1, 1]
        #         J[i * 2 + 1, 4] = -X * Y / Z_2 * self.K[1, 1]
        #         J[i * 2 + 1, 5] = -X / Z * self.K[1, 1]
        #     return J

        # P_ba = bundle_adjustment(self.K, measurements, database, P_ba)
        result = least_squares(
            residual_function, initial_pose, args=(measurements, database)
        )
        if not result.success:
            return False, None, None
        t = result.x[3:][:, None]  # 提取平移部分
        R = cv2.Rodrigues(result.x[:3])[0]  # 提取旋转部分
        return True, R, t
